#pragma once

#include <iostream>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <vector>
#include <functional>
#include <future>
#include <atomic>

#define DEFAULT_N_THREAD 1

namespace mq {

    class ThreadPool {
    public:
        using ThreadFn = std::function<void(void)>;

        ThreadPool(int n_threads = DEFAULT_N_THREAD) 
            : _stop(false)
        {
            for (int i = 0; i < n_threads; i ++)
                _thr_pool.emplace_back(&ThreadPool::thread_loop, this);
        }

        /* 
        *  不知道任务的类型，也不知道任务的参数, 也不知道任务的返回值
        *  需要对 packaged_task 使用 shared_ptr 包装，不然出了作用域就会被销毁
        *  
        *  统一封装所有用户任务, 并扔到任务队列中
        *  先 绑定一下用户任务参数，并且使用完美转发保证参数的类型
        *  使用 packaged_task 对该 bind 继续再次封装
        */ 
        template<typename Fn, typename ... Args>
        auto submit(Fn && fn, Args && ... args) -> std::future<decltype(fn(args ...))> {
            using ReturnType = decltype(fn(args ...));

            auto fn_ = std::bind(std::forward<Fn>(fn), std::forward<Args>(args) ...);
            auto task = std::make_shared<std::packaged_task<ReturnType()>> (fn_);

            std::future<ReturnType> fut = task->get_future();

            // 往任务队列中添加新任务
            {
                std::unique_lock<std::mutex> lock(_mtx);
                _taskq.push_back([task] () {
                    (*task)();
                });
                _cv.notify_one();
            }

            return fut;
        }

        // 唤醒所有线程并回收所有线程
        void stop() {
            _stop = true;
            _cv.notify_all();
            for (auto& thr : _thr_pool) {
                if (thr.joinable()) {
                    thr.join();
                }
            }
        }

        ~ThreadPool()
        {
            stop();
        }
    private:
        // 每个线程的任务
        void thread_loop() {
            while (!_stop) {
                /*
                * 为了避免频繁加锁，这里一次取出任务队列中的所有任务
                */
                std::vector<ThreadFn> taskq;
                {
                    // 等待到任务池不为空
                    std::unique_lock<std::mutex> lock(_mtx);
                    _cv.wait(lock, [this] () {
                        return !_taskq.empty() || _stop;
                    });

                    taskq.swap(_taskq);
                }

                for (auto& task : taskq) {
                    task();
                }
            }
        }

    private:
        std::mutex                  _mtx;
        std::condition_variable     _cv;
        std::atomic<bool>           _stop;
        std::vector<ThreadFn>       _taskq;
        std::vector<std::thread>    _thr_pool;
    };

    using ThreadPoolPtr = std::shared_ptr<ThreadPool>;
}